#!/usr/bin/env python
# _*_ coding:utf-8 _*_
__author__ = 'ybh'

import socket
import threading
import sys
import os
import base64
import hashlib
import struct
import time
import array
import logging
from logging.handlers import TimedRotatingFileHandler

# ====== config ======
HOST = '0.0.0.0'
PORT = 10086
MAGIC_STRING = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
HANDSHAKE_STRING = "HTTP/1.1 101 Switching Protocols\r\n" \
                   "Upgrade:websocket\r\n" \
                   "Connection: Upgrade\r\n" \
                   "Sec-WebSocket-Accept: {1}\r\n" \
                   "WebSocket-Location: ws://{2}/chat\r\n" \
                   "WebSocket-Protocol:chat\r\n\r\n"

LOG_FILE="/var/log/ws_server.log"
logger_name='ws_server'
logger=logging.getLogger(logger_name)
logger.setLevel(logging.INFO)

formatter = logging.Formatter('%(asctime)s %(levelname)-8s: %(message)s')
fh = logging.FileHandler(LOG_FILE)
fh.setFormatter(formatter)
fh.setLevel(logging.INFO)
logger.addHandler(fh)

#添加日志切割
hdlr = TimedRotatingFileHandler(LOG_FILE,when='d',interval=1,backupCount=3)
logger.addHandler(hdlr)



class Th(threading.Thread):
    def __init__(self, connection, ):
        threading.Thread.__init__(self)
        self.con = connection

    def run(self):
        try:
            message = self.recv_data(4096)
            message = eval(message)
            file_name = message['file_name']
            size = message['size']
            path = message['path']
            if not os.path.exists(path):
                os.makedirs(path)

            file_path = "%s/%s" % (path, file_name)
            f = open(file_path, 'wb+')
            while True:
                self.send_data('ok')
                slice = self.receive_frame(self.con)
                if slice[2] == 'end' or slice[2] == '':
                    break
                else:
                    f.write(slice[2])
                    time.sleep(0.02)
            f.close()
            self.con.close()
        except Exception as e:
            logger.exception(u"传输异常")
        self.con.close()

    def recv_data(self, num):
        try:
            all_data = self.con.recv(num)
            if not len(all_data):
                return False
        except:
            return False
        else:
            code_len = ord(all_data[1]) & 127
            if code_len == 126:
                masks = all_data[4:8]
                data = all_data[8:]
            elif code_len == 127:
                masks = all_data[10:14]
                data = all_data[14:]
            else:
                masks = all_data[2:6]
                data = all_data[6:]
            raw_str = ""
            i = 0
            for d in data:
                raw_str += chr(ord(d) ^ ord(masks[i % 4]))
                i += 1
            return raw_str

    def receive_frame(self, sock):
        bytes_1_2 = sock.recv(2)
        byte_1 = ord(bytes_1_2[0])
        byte_2 = ord(bytes_1_2[1])
        # 读最高位,结束标识
        is_finished = byte_1 >> 7 & 1
        # 低4位是传输数据的类型
        opcode = byte_1 & 0xf
        if opcode == 0x8:
            logger.error('connection closed.')
            return is_finished, opcode, ''
        # 第2个字节的最高位表示是否对数据进行了掩码运算
        is_mask = byte_2 >> 7 & 1
        # 低7位表示传输数据的长度
        length = byte_2 & 0x7f

        length_data = ""
        # 如果负载长度是126,之后的两个字节也将被解释为16位无符号整数的负载长度
        if length == 0x7e:
            length_data = sock.recv(2)
            length = struct.unpack("!H", length_data)[0]
        # 如果负载长度是127,之后的8个字节也将被解释为一个64位无符号整数的负载长度
        elif length == 0x7f:
            length_data = sock.recv(8)
            length = struct.unpack("!Q", length_data)[0]
        mask_key = ""
        # 如果进行了掩码运算,要读出32位的掩码数据
        if is_mask:
            mask_key = sock.recv(4)
        # 接下来就是真正传输的数据了
        data = sock.recv(length)
        if is_mask:
            data = self.mask_or_unmask(mask_key, data)
        return is_finished, opcode, data

    def mask_or_unmask(self, mask_key, data):
        # 1个字节的int
        _m = array.array("B", mask_key)
        _d = array.array("B", data)
        # mask_key是4个字节,所以余4
        for i in xrange(len(_d)):
            _d[i] ^= _m[i % 4]
        return _d.tostring()
        # send data

    def send_data(self, data):
        if data:
            data = str(data)
        else:
            return False
        token = "\x81"
        length = len(data)
        if length < 126:
            token += struct.pack("B", length)
        elif length <= 0xFFFF:
            token += struct.pack("!BH", 126, length)
        else:
            token += struct.pack("!BQ", 127, length)
        # struct为Python中处理二进制数的模块，二进制流为C，或网络流的形式。
        data = '%s%s' % (token, data)
        self.con.send(data)
        return True


        # handshake


def handshake(con):
    headers = {}
    shake = con.recv(1024)

    if not len(shake):
        return False

    header, data = shake.split('\r\n\r\n', 1)
    for line in header.split('\r\n')[1:]:
        key, val = line.split(': ', 1)
        headers[key] = val

    if 'Sec-WebSocket-Key' not in headers:
        logger.error('This socket is not websocket, client close.')
        con.close()
        return False

    sec_key = headers['Sec-WebSocket-Key']
    res_key = base64.b64encode(hashlib.sha1(sec_key + MAGIC_STRING).digest())

    str_handshake = HANDSHAKE_STRING.replace('{1}', res_key).replace('{2}', HOST + ':' + str(PORT))
    logger.info(str_handshake)
    con.send(str_handshake)
    return True


def new_service():
    """start a service socket and listen
    when coms a connection, start a new thread to handle it"""

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    try:
        sock.bind((HOST, PORT))
        sock.listen(1000)
        # 链接队列大小
        logger.info("bind %s,ready to use" % PORT)
    except Exception as e:
        logger.exception("Server is already running,quit")
        sys.exit()

    while True:
        connection, address = sock.accept()
        # 返回元组（socket,add），accept调用时会进入waite状态
        logger.info("Got connection from %s %s" % (address[0],address[1]))
        if handshake(connection):
            logger.info("handshake success")
            try:
                t = Th(connection)
                t.start()
                logger.info('new thread for client ...')
            except Exception as e:
                logger.exception('start new thread error')
                connection.close()


if __name__ == '__main__':
    new_service()
